import numpy as np
import matplotlib.pyplot as plt
from ipywidgets import interact, FloatSlider
import warnings

# 抑制非关键数值警告（仅调试时使用，正式报告需移除）
warnings.filterwarnings("ignore", category=RuntimeWarning, message="invalid value encountered in divide")
warnings.filterwarnings("ignore", category=RuntimeWarning, message="invalid value encountered in log")

plt.rcParams['font.sans-serif'] = ['SimHei']  # 中文显示
plt.rcParams['axes.unicode_minus'] = False    # 负号显示
plt.rcParams['figure.dpi'] = 300             # 高分辨率输出


def complex_potential(z, U, Gamma, a, eps=1e-8):
    """复势函数计算（严格复数安全处理）"""
    # 处理z=0和负实数情况（确保z_safe为非零复数）
    z_real = np.where(np.abs(z.real) < eps, eps, z.real)
    z_imag = np.where(np.abs(z.imag) < eps, eps, z.imag)
    z_safe = z_real + 1j * z_imag  # 分别保护实部和虚部
    
    return U * (z_safe + a**2 / z_safe) + 1j * Gamma / (2 * np.pi) * np.log(z_safe)


def complex_velocity(z, U, Gamma, a, eps=1e-8):
    """复速度场计算（严格数值保护）"""
    z_real = np.where(np.abs(z.real) < eps, eps, z.real)
    z_imag = np.where(np.abs(z.imag) < eps, eps, z.imag)
    z_safe = z_real + 1j * z_imag  # 实虚部分离保护
    
    return U * (1 - a**2 / z_safe**2) + 1j * Gamma / (2 * np.pi * z_safe)


def stagnation_points_analytic(U, Gamma, a):
    """解析法求解驻点（适用于含单涡的圆柱绕流）"""
    Gamma_max = 4 * np.pi * U * a
    if abs(Gamma) > Gamma_max:
        raise ValueError(f"环量Γ={Gamma}超过安全范围Γ∈[-{Gamma_max:.1f}, {Gamma_max:.1f}]")
    
    term = Gamma / (4 * np.pi * U)
    sqrt_term = np.sqrt(a**2 - term**2)  # 实数解（安全环量内）
    z1 = 1j * term + sqrt_term  # 上驻点（Γ>0时下移，Γ<0时上移）
    z2 = 1j * term - sqrt_term  # 下驻点
    return np.array([z1, z2])


def plot_flow_field(U=1.0, Gamma=5.0, a=1.0):
    """固定xlim和ylim为元组，避免交互参数类型错误"""
    xlim = (-5, 5)
    ylim = (-5, 5)
    
    # 生成计算网格（排除圆柱内部）
    x = np.linspace(xlim[0], xlim[1], 300)
    y = np.linspace(ylim[0], ylim[1], 300)
    X, Y = np.meshgrid(x, y)
    Z = X + 1j * Y
    Z[np.abs(Z) < a] = np.nan  # 屏蔽圆柱内部区域

    # 计算物理速度场（复速度共轭关系，严格数值保护）
    V = complex_velocity(Z, U, Gamma, a)
    u_phys = V.real    # 物理速度u分量（x方向）
    v_phys = -V.imag   # 物理速度v分量（y方向，复速度定义为u - iv）

    # 求解并筛选有效驻点（位于圆柱外）
    try:
        z_stag = stagnation_points_analytic(U, Gamma, a)
        z_stag = z_stag[np.abs(z_stag) >= a - 1e-6]  # 允许数值误差范围内的边界驻点
    except ValueError as e:
        z_stag = []
        plt.text(0.5, 0.95, f"错误：{e}", transform=plt.gca().transAxes,
                 color='red', fontsize=12, bbox=dict(facecolor='white', alpha=0.8))

    # 初始化画布
    fig, ax = plt.subplots(figsize=(10, 8))

    # 绘制流线（动态流场，移除不兼容的alpha参数）
    ax.streamplot(X, Y, u_phys, v_phys, color='blue', density=1.5,
                  linewidth=0.8, arrowsize=1.2, zorder=1)

    # 绘制等势线（静态等值线）
    ax.contour(X, Y, complex_potential(Z, U, Gamma, a).real,
               levels=15, colors='red', linestyles='dashed',
               linewidths=0.6, zorder=2)

    # 标记驻点
    for z in z_stag:
        ax.scatter(z.real, z.imag, s=150, c='gold', edgecolor='black',
                   marker='*', label='驻点' if len(z_stag) == 1 else "", zorder=3)

    # 绘制圆柱边界
    theta = np.linspace(0, 2 * np.pi, 200)
    ax.plot(a * np.cos(theta), a * np.sin(theta), 'k-', lw=2, label='圆柱表面')

    # 安全环量提示
    Gamma_max = 4 * np.pi * U * a
    if abs(Gamma) > Gamma_max:
        ax.text(0.5, 0.95, f"警告：Γ={Gamma:.1f} > Γ_max={Gamma_max:.1f}",
                ha='center', va='top', transform=ax.transAxes,
                color='red', fontsize=12, bbox=dict(facecolor='white', alpha=0.8))

    # 图例与标注
    ax.set_title(f"环量调控流场 (U={U}, Γ={Gamma:.1f}, a={a})", fontsize=14)
    ax.set_xlabel('x', fontsize=12), ax.set_ylabel('y', fontsize=12)
    ax.legend(loc='upper right', fontsize=10)
    ax.grid(True, linestyle='--', alpha=0.6)
    ax.axis('equal'), ax.set_xlim(xlim), ax.set_ylim(ylim)

    plt.tight_layout()
    plt.show()


def verify_circulation(Gamma=5.0, a=1.0, N=2000):
    """沿圆柱表面积分验证环量守恒"""
    theta = np.linspace(0, 2 * np.pi, N, endpoint=False)
    z = a * np.exp(1j * theta)  # 圆柱表面离散点
    dz_analytic = 1j * a * np.exp(1j * theta) * (2 * np.pi / N)  # 解析弧长元素

    # 计算复速度并积分（严格数值保护）
    V = complex_velocity(z, U=1.0, Gamma=Gamma, a=a)
    circulation_num = np.sum(V * dz_analytic).real  # 数值积分结果
    circulation_theory = -Gamma  # 理论环量（复势定义中的符号关系）

    # 误差分析
    relative_error = abs((circulation_num - circulation_theory) / circulation_theory) * 100

    # 结果输出
    print(f"■ 环量验证结果 (Γ={Gamma}, a={a})")
    print(f"· 理论值：Γ_theory = {circulation_theory:.4f}")
    print(f"· 数值积分值：Γ_num = {circulation_num:.4f}")
    print(f"· 相对误差：{relative_error:.2f}% （目标：≤0.5%）\n")

    # 可视化积分路径
    fig, ax = plt.subplots(figsize=(6, 6))
    ax.plot(z.real, z.imag, 'b-', lw=2, label='积分路径')
    ax.scatter(z.real[0], z.imag[0], c='red', marker='o', label='起点')
    ax.set_title("环量积分路径（圆柱表面离散点）", fontsize=12)
    ax.set_xlabel('x'), ax.set_ylabel('y')
    ax.axis('equal'), ax.legend()
    ax.grid(True, linestyle='--', alpha=0.5)
    plt.show()


# 交互式流场可视化（仅对U、Gamma、a添加滑块，xlim/ylim固定为元组）
interact(
    plot_flow_field,
    U=FloatSlider(min=0.5, max=3.0, step=0.1, value=1.0, description='来流速度 U'),
    Gamma=FloatSlider(min=-15, max=15, step=0.5, value=0.0, description='环量 Γ'),
    a=FloatSlider(min=0.5, max=2.0, step=0.1, value=1.0, description='圆柱半径 a')
);

# 交互式环量验证（同理，固定xlim/ylim）
interact(
    verify_circulation,
    Gamma=FloatSlider(min=-10, max=10, step=0.5, value=5.0, description='环量 Γ'),
    a=FloatSlider(min=0.5, max=2.0, step=0.1, value=1.0, description='圆柱半径 a')
);